import configparser
import multiprocessing
from datetime import datetime, timedelta
from joblib import Parallel, delayed
from tqdm import tqdm
from ..common.utils import collection_util
from ..common.db import stock_info

class StockLabelGenerator:
    def __init__(self):
        self.quant_config = configparser.ConfigParser()
        self.quant_config.read('../config/quant_config.ini')

    def init_tables(self):
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        for i in range(20):
            cursor.execute(stock_label_create_sql.format(i))
        db.close()

    def build_or_update_label_20d(self, biz_date):
        """
        构建/更新股票后20日收益回撤label
          - 用code hash取模做分表条件
          - 如果后20个交易日价格均低于当日收盘价，那么无收益，最大收益取0
          - 如果后20个交易日价格均高于当日收盘价，那么无回撤，最大回撤取0
        """
        # self.__execute_update_label_20d_per_stock('600900', biz_date)
        code_set = stock_info.get_all_stock_code_set()
        Parallel(n_jobs=multiprocessing.cpu_count())(delayed(StockLabelGenerator.__execute_update_label_20d_per_stock)(
            self, code, biz_date) for code in tqdm(code_set, '更新股票20d label流水表'))

    def __execute_update_label_20d_per_stock(self, code, biz_date):
        earliest_date = self.quant_config['stock_info']['earliest_date']
        table_name = stock_info.get_stock_label_day_table_name(code)
        table_name_day_k = stock_info.get_stock_hist_day_k_table_name(code)
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        cursor.execute('SELECT max(`date`) FROM {} WHERE `code` = %s'.
                       format(table_name), code)
        recent_row = cursor.fetchone()
        start_date = (recent_row[0] + timedelta(days=1)).strftime('%Y%m%d') if recent_row is not None and recent_row[0] is not None else earliest_date
        cursor.execute('SELECT `date`, `close_hfq_by_factor`, `high_hfq_by_factor`, `low_hfq_by_factor` FROM {} WHERE `code` = %s AND (`date` BETWEEN %s and %s)'.format(table_name_day_k), [code, start_date, biz_date])
        data = cursor.fetchall()
        if len(data) == 0:
            return
        date_list = list()
        close_list = list()
        high_hfq_by_factor_list = list()
        low_hfq_by_factor_list = list()
        for row in data:
            date_list.append(row[0])
            close_list.append(row[1])
            high_hfq_by_factor_list.append(row[2])
            low_hfq_by_factor_list.append(row[3])
        high_hfq_by_factor_slide_list = collection_util.max_sliding_window(high_hfq_by_factor_list, 20)
        low_hfq_by_factor_slide_list = collection_util.min_sliding_window(low_hfq_by_factor_list, 20)
        label_info_list = list()
        for i in range(len(high_hfq_by_factor_list)):
            if i + 1 >= len(high_hfq_by_factor_slide_list):
                break
            high_hfq_by_factor_in_20d = high_hfq_by_factor_slide_list[i + 1]
            low_hfq_by_factor_in_20d = low_hfq_by_factor_slide_list[i + 1]
            date = date_list[i]
            max_increase_20d = max((high_hfq_by_factor_in_20d - close_list[i])/close_list[i], 0)
            max_decrease_20d = max((close_list[i] - low_hfq_by_factor_in_20d)/close_list[i], 0)
            label_info_list.append((code, date, max_increase_20d, max_decrease_20d))
        if len(label_info_list) > 0:
            sql = (('INSERT INTO {} (`code`, `date`, `max_increase_20d`, `max_decrease_20d`, `max_increase_50d`, '
                    '`max_decrease_50d`, `max_increase_120d`, `max_decrease_120d`) '
                    'VALUES (%s, %s, %s, %s, NULL, NULL, NULL, NULL)'
                    ).format(
                table_name))
            cursor.executemany(sql, label_info_list)
            db.commit()
        cursor.close()
        db.close()

    def build_or_update_label_50d(self, biz_date):
        """
        构建/更新股票后50日收益回撤label
          - 用code hash取模做分表条件
          - 如果后50个交易日价格均低于当日收盘价，那么无收益，最大收益取0
          - 如果后50个交易日价格均高于当日收盘价，那么无回撤，最大回撤取0
        """
        # self.__execute_update_label_50d_per_stock('600900', biz_date)
        code_set = stock_info.get_all_stock_code_set()
        Parallel(n_jobs=multiprocessing.cpu_count())(delayed(StockLabelGenerator.__execute_update_label_50d_per_stock)(
            self, code, biz_date) for code in tqdm(code_set, '更新股票50d label流水表'))

    def __execute_update_label_50d_per_stock(self, code, biz_date):
        earliest_date = self.quant_config['stock_info']['earliest_date']
        table_name = stock_info.get_stock_label_day_table_name(code)
        table_name_day_k = stock_info.get_stock_hist_day_k_table_name(code)
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        cursor.execute('SELECT max(`date`) FROM {} WHERE `code` = %s AND `max_increase_50d` is not NULL AND `max_decrease_50d` is not NULL'.
                       format(table_name), code)
        recent_row = cursor.fetchone()
        start_date = (recent_row[0] + timedelta(days=1)).strftime('%Y%m%d') if recent_row is not None and recent_row[0] is not None else earliest_date
        cursor.execute('SELECT `date`, `close_hfq_by_factor`, `high_hfq_by_factor`, `low_hfq_by_factor` FROM {} WHERE `code` = %s AND (`date` BETWEEN %s and %s)'.format(table_name_day_k), [code, start_date, biz_date])
        data = cursor.fetchall()
        if len(data) == 0:
            return
        date_list = list()
        close_list = list()
        high_hfq_by_factor_list = list()
        low_hfq_by_factor_list = list()
        for row in data:
            date_list.append(row[0])
            close_list.append(row[1])
            high_hfq_by_factor_list.append(row[2])
            low_hfq_by_factor_list.append(row[3])
        high_hfq_by_factor_slide_list = collection_util.max_sliding_window(high_hfq_by_factor_list, 50)
        low_hfq_by_factor_slide_list = collection_util.min_sliding_window(low_hfq_by_factor_list, 50)
        label_info_list = list()
        for i in range(len(high_hfq_by_factor_list)):
            if i + 1 >= len(high_hfq_by_factor_slide_list):
                break
            high_hfq_by_factor_in_50d = high_hfq_by_factor_slide_list[i + 1]
            low_hfq_by_factor_in_50d = low_hfq_by_factor_slide_list[i + 1]
            date = date_list[i]
            max_increase_50d = max((high_hfq_by_factor_in_50d - close_list[i])/close_list[i], 0)
            max_decrease_50d = max((close_list[i] - low_hfq_by_factor_in_50d)/close_list[i], 0)
            label_info_list.append((max_increase_50d, max_decrease_50d, code, date))
        if len(label_info_list) > 0:
            sql = (('UPDATE {} SET `max_increase_50d` = %s, `max_decrease_50d` = %s '
                    'WHERE `code`=%s AND `date`=%s'
                    ).format(
                table_name))
            cursor.executemany(sql, label_info_list)
            db.commit()
        cursor.close()
        db.close()

    def build_or_update_label_120d(self, biz_date):
        """
        构建/更新股票后120日收益回撤label
          - 用code hash取模做分表条件
          - 如果后120个交易日价格均低于当日收盘价，那么无收益，最大收益取0
          - 如果后120个交易日价格均高于当日收盘价，那么无回撤，最大回撤取0
        """
        # self.__execute_update_label_120d_per_stock('600900', biz_date)
        code_set = stock_info.get_all_stock_code_set()
        Parallel(n_jobs=multiprocessing.cpu_count())(delayed(StockLabelGenerator.__execute_update_label_120d_per_stock)(
            self, code, biz_date) for code in tqdm(code_set, '更新股票120d label流水表'))

    def __execute_update_label_120d_per_stock(self, code, biz_date):
        earliest_date = self.quant_config['stock_info']['earliest_date']
        table_name = stock_info.get_stock_label_day_table_name(code)
        table_name_day_k = stock_info.get_stock_hist_day_k_table_name(code)
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        cursor.execute('SELECT max(`date`) FROM {} WHERE `code` = %s AND `max_increase_120d` is not NULL AND `max_decrease_120d` is not NULL'.
                       format(table_name), code)
        recent_row = cursor.fetchone()
        start_date = (recent_row[0] + timedelta(days=1)).strftime('%Y%m%d') if recent_row is not None and recent_row[0] is not None else earliest_date
        cursor.execute('SELECT `date`, `close_hfq_by_factor`, `high_hfq_by_factor`, `low_hfq_by_factor` FROM {} WHERE `code` = %s AND (`date` BETWEEN %s and %s)'.format(table_name_day_k), [code, start_date, biz_date])
        data = cursor.fetchall()
        if len(data) == 0:
            return
        date_list = list()
        close_list = list()
        high_hfq_by_factor_list = list()
        low_hfq_by_factor_list = list()
        for row in data:
            date_list.append(row[0])
            close_list.append(row[1])
            high_hfq_by_factor_list.append(row[2])
            low_hfq_by_factor_list.append(row[3])
        high_hfq_by_factor_slide_list = collection_util.max_sliding_window(high_hfq_by_factor_list, 120)
        low_hfq_by_factor_slide_list = collection_util.min_sliding_window(low_hfq_by_factor_list, 120)
        label_info_list = list()
        for i in range(len(high_hfq_by_factor_list)):
            if i + 1 >= len(high_hfq_by_factor_slide_list):
                break
            high_hfq_by_factor_in_120d = high_hfq_by_factor_slide_list[i + 1]
            low_hfq_by_factor_in_120d = low_hfq_by_factor_slide_list[i + 1]
            date = date_list[i]
            max_increase_120d = max((high_hfq_by_factor_in_120d - close_list[i])/close_list[i], 0)
            max_decrease_120d = max((close_list[i] - low_hfq_by_factor_in_120d)/close_list[i], 0)
            label_info_list.append((max_increase_120d, max_decrease_120d, code, date))
        if len(label_info_list) > 0:
            sql = (('UPDATE {} SET `max_increase_120d` = %s, `max_decrease_120d` = %s '
                    'WHERE `code`=%s AND `date`=%s'
                    ).format(
                table_name))
            cursor.executemany(sql, label_info_list)
            db.commit()
        cursor.close()
        db.close()

    def clear_all_data(self):
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        for i in range(20):
            cursor.execute('DROP TABLE IF EXISTS stock_label_day_{}'.format(str(i)))
            db.commit()
        cursor.close()
        db.close()


stock_label_create_sql = '''
create table if not exists `stock_label_day_{}` (
    `code` VARCHAR(20) NOT NULL COMMENT '股票代码',
    `date` TIMESTAMP NOT NULL COMMENT '交易日期',
    `max_increase_20d` DOUBLE NOT NULL COMMENT '后20日最大涨幅（如果后续交易日收盘价均低于当日收盘价，该值为0）-后复权',
    `max_decrease_20d` DOUBLE NOT NULL COMMENT '后20日最大回撤（如果后续交易日收盘价均高于当日收盘价，该值为0）-后复权',
    `max_increase_50d` DOUBLE COMMENT '后50日最大涨幅（如果后续交易日收盘价均低于当日收盘价，该值为0）-后复权',
    `max_decrease_50d` DOUBLE COMMENT '后50日最大回撤（如果后续交易日收盘价均高于当日收盘价，该值为0）-后复权',
    `max_increase_120d` DOUBLE COMMENT '后50日最大涨幅（如果后续交易日收盘价均低于当日收盘价，该值为0）-后复权',
    `max_decrease_120d` DOUBLE COMMENT '后50日最大回撤（如果后续交易日收盘价均高于当日收盘价，该值为0）-后复权',
    KEY (`code`),
    PRIMARY KEY (`code`, `date`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8 COMMENT '股票样本label流水表';
'''